"""
@Description :   MLLM object detection 能力测试
@Author      :   tqychy 
@Time        :   2025/08/28 08:46:41
"""
import sys

sys.path.append("./testers")
import json
import os

import matplotlib
import matplotlib.pyplot as plt
from base_tester import BaseTester
from matplotlib.patches import Rectangle
from tqdm import tqdm

matplotlib.use('Agg')


class DetectTester(BaseTester):

    @staticmethod
    def add_bboxes(image, bboxes, categories, bbox_color, image_save_path, cat_color="purple"):
        """
        为图像添加 bboxes 和标签
        """
        if bboxes is None or len(bboxes) == 0:
            return
        ax = plt.gca()
        ax.imshow(image)
        for category, bbox in zip(categories, bboxes):
            if bbox is None:
                continue
            box_plot = Rectangle(
                (bbox[0], bbox[1]), bbox[2], bbox[3], fill=False, edgecolor=bbox_color, linewidth=3)
            ax.add_patch(box_plot)
            ax.text(bbox[0], bbox[1] - 10, category, fontsize=8, color=cat_color, bbox=dict(facecolor='white', alpha=0.7, pad=2))
        
        plt.savefig(image_save_path)
        plt.clf()

    def parse_output(self, output):
        model_name = self.cfg.EVAL.MODEL_NAME
        if model_name == "qwen2.5-VL" or model_name == "qwen2.5-VL-API" or model_name == "gemini":
            lines = output.splitlines()
            for i, line in enumerate(lines):
                if line == "```json":
                    json_output = "\n".join(lines[i+1:])  # Remove everything before "```json"
                    json_output = json_output.split("```")[0]  # Remove everything after the closing "```"
                    break  # Exit the loop once "```json" is found
            outputs = json.loads(json_output)
            for output in outputs:
                if "bbox_2d" in output.keys():
                    output["bbox"] = output["bbox_2d"]
                    del output["bbox_2d"]
                if "label" in output.keys():
                    output["category"] = output["label"]
                    del output["label"]
                score = output.get("score", 1.0)
                output["score"] = score if score is not None else 1.0
            return outputs
        elif model_name == "glm-4.5V":
            start_idx = output.find("<|begin_of_box|>")
            end_idx = output.find("<|end_of_box|>")
            if start_idx == -1 or end_idx == -1:
                raise ValueError(f"bbox tag not found! {output}")
            json_content = output[start_idx + len("<|begin_of_box|>"):end_idx]
            outputs = json.loads(json_content)
            for output in outputs:
                output["score"] = output.get("score", 1.0)
            return outputs
        elif model_name == "dino-x" or model_name == "dino-xseek":
            return output
        elif model_name == "grounding-dino-base":
            bboxes = output["boxes"].cpu()
            categories = output["text_labels"]
            scores = output["scores"].cpu()
            outputs = []
            for bbox, (category, score) in zip(bboxes, zip(categories, scores)):
                if len(category) == 0 or len(bbox) != 4:
                    continue
                outputs.append({
                    "bbox": bbox.tolist(),
                    "category": category,
                    "score": score.item()
                })
            return outputs
        else:
            raise ValueError(f"unknown model name {model_name}")
        
    @staticmethod
    def ensure_class(pred_category, gt_categories):
        def lcs(s1, s2):
            m, n = len(s1), len(s2)
            max_len = 0
            dp = [[0] * (n + 1) for _ in range(m + 1)]
            for i in range(1, m + 1):
                for j in range(1, n + 1):
                    if s1[i-1] == s2[j-1]:
                        dp[i][j] = dp[i-1][j-1] + 1
                        max_len = max(max_len, dp[i][j])
                    else:
                        dp[i][j] = 0
            return max_len
        categories = list(set(gt_categories))
        pasred_gt_categories = [cat.lower().strip().replace(" ", "").replace(".", "").replace("#", "") for cat in categories]
        parsed_pred_category = pred_category.lower().strip().replace(" ", "").replace(".", "").replace("#", "")
        try:
            ind = pasred_gt_categories.index(parsed_pred_category)
        except:
            lcs_score = [lcs(parsed_pred_category, cat) for cat in pasred_gt_categories]
            max_len = max(lcs_score)
            ind = lcs_score.index(max_len)
        return categories[ind]
    
    def test(self):
        with tqdm(total=len(self.dataset)) as pbar:
            for data in self.dataloader:
                image_path = data["image_path"][0]
                image_name = os.path.basename(image_path)
                image = self.load_image(image_path)
                gt_categories = [cat[0] for cat in data["category"]]
                gt_bboxes = data["bbox"]

                inputs = self.make_prompt(image_path, "", ".".join(list(set(gt_categories))))
                try:
                    outputs = self.model(inputs)
                    outputs = self.parse_output(outputs)
                except Exception as e:
                    self.logger.error(f"img {image_name} failed with {e}: model output error!")
                    map50, f1 = self.metrics_handler.desc_metrics()
                    pbar.desc = f"mAP50/f1: {map50:.3f}/{f1:.3f}"
                    pbar.update(1)
                    continue
                pred_bboxes = []
                pred_categories = []
                scores = []
                for output in outputs:
                    try:
                        bbox = output["bbox"]
                        bbox = self.covert_formatted_bbox(bbox, image.shape)
                        pred_bboxes.append(bbox)
                        pred_categories.append(self.ensure_class(output["category"], gt_categories))
                        scores.append(output["score"])
                    except Exception as e:
                        self.logger.warning(f"part of the output of img {image_name} failed with {e}: {output}")
                        continue
                
                self.metrics_handler.update(pred_bboxes, pred_categories, scores, gt_bboxes, gt_categories)
                self.logger.info(f"img {image_name}: {self.metrics_handler.display()}")
                self.add_bboxes(image, pred_bboxes, pred_categories, "orange", os.path.join(self.pred_save_path, image_name))
                self.add_bboxes(image, gt_bboxes, gt_categories, "green", os.path.join(self.gt_save_path, image_name))
                map50, f1 = self.metrics_handler.desc_metrics()
                pbar.desc = f"mAP50/f1: {map50:.3f}/{f1:.3f}"
                pbar.update(1)
        
        self.logger.critical(self.metrics_handler.display_all())
